#!/usr/bin/env python
# -*- coding: utf-8 -*-
# @Time    : 2023/4/23 11:01
# @Author  : WJQ
# @Site    :
# @File    : mitmproxy_handler.py
# @Software: PyCharm
import time
from mitmproxy import http, ctx
import sys, os

# 获取当前脚本的绝对路径
current_script_path = os.path.abspath(__file__)

# 获取上层目录的上层目录的路径
grandparent_dir = os.path.dirname(current_script_path)

# 将这个路径添加到sys.path
sys.path.append(grandparent_dir)


# from attack.get_suspicious_request import TrafficReplay
from database.mongo_handler import MongoHandler
import sys

# logging.basicConfig(filename='./app.log', filemode='w', format='%(name)s - %(levelname)s - %(message)s')
# logger = logging.getLogger(__name__)

# 获取传递给脚本的 file_name 参数
# 初始化默认值
file_name = "test-10.23"

# 通过遍历 sys.argv 获取 file_name 参数
for index, arg in enumerate(sys.argv):
    if arg == "--file-name" and index < len(sys.argv) - 1:
        file_name = sys.argv[index + 1]
        break


def write_to_file(dir, data):
    global file_name
    file_name_txt = './{}/{}.txt'.format(dir, file_name)
    # 'a+'表示如果文件不存在，就创建该文件。然后以追加模式打开该文件。
    # 追加模式是以写的方式打开一个文件用于在末尾添加数据。
    with open(file_name_txt, 'a+') as file:
        file.write(data + '\n')  # 添加数据，并在数据后面添加一个换行符。


class MitmProxyHandler:
    def __init__(self, applet_name, pageUrl):
        self.applet_name = applet_name
        self.pageUrl = pageUrl
        # dbName
        self.mongo = MongoHandler("traffic")
        # self.traffic_replay = TrafficReplay()

    def request(self, flow: http.HTTPFlow) -> None:
        if flow.request.pretty_url.endswith(('.jpg', '.jpeg', '.png', '.gif', '.css', '.js')):
            flow.request.anticache()
            # flow.response = http.HTTPResponse.make(200)
        flow.request.pageUrl = self.pageUrl.value

    def is_static_resource(self, url):
        static_file_extensions = [".jpg", ".jpeg", ".png", ".gif", ".css", ".js", ".woff", ".woff2", ".ttf", ".svg", ".ico", ".json"]
        for ext in static_file_extensions:
            if url.endswith(ext):
                return True
        return False

    def is_mini_program_traffic(self, headers):
        user_agent = headers.get("User-Agent")
        return ("MiniProgram" in user_agent or "miniProgram" in user_agent)

    def response(self, flow: http.HTTPFlow) -> None:
        if flow.request.pretty_host and self.is_mini_program_traffic(flow.request.headers):
            status_code = flow.response.status_code
            if status_code // 100 == 5:
                ctx.log.info(f"Discarding traffic with status code 5xx: {flow.request.url}")
                return

            if status_code == 204:
                ctx.log.info(f"Skipping traffic with status code 204: {flow.request.url}")
                return

            if self.is_static_resource(flow.request.url):
                ctx.log.info(f"Skipping static resource: {flow.request.url}")
                return

            ctx.log.info(f"Capturing traffic: {flow.request.url}")

            req_data = {
                "url": flow.request.url,
                "method": flow.request.method,
                "headers": dict(flow.request.headers),
                "content": flow.request.get_text(strict=False),
                "url_form": flow.request.urlencoded_form,
                "timestamp_start": flow.request.timestamp_start,
                "page_url": flow.request.pageUrl
            }

            res_data = {
                "status_code": status_code,
                "headers": dict(flow.response.headers),
                "content": flow.response.get_text(strict=False)
            }

            data = {
                "request": req_data,
                "response": res_data
            }

            if not self.mongo.is_data_exists(self.applet_name.value, data):
                # logger.info("find a flow, insert data")
                print(r"[+] {} : find a flow, insert data".format(time.strftime("%Y-%m-%d %H:%M:%S", time.localtime())))
                self.mongo.insert_data(self.applet_name.value, data)
                # write_to_file('record_all', str(data['_id']) + "---" + str(data['request']['url']))
                # if self.traffic_replay.judge_info_sensitive(data):
                #     print(r"[+] {} : suspicious flow {}".format(time.strftime("%Y-%m-%d %H:%M:%S", time.localtime()), data['request']['url']))
                #     write_to_file('record', str(data['_id']) + "---" + str(data['request']['url']))
            else:
                print(r"[-] {} : find a flow, but {} exists".format(time.strftime("%Y-%m-%d %H:%M:%S", time.localtime()), data['request']['url']))
        else:
            ctx.log.info(f"Skipping non-WeChat traffic: {flow.request.url}")


# print("now is database : {}".format(file_name))

addons = [
    # MitmProxyHandler(storage_folder="experiment-statistics")
]

# mitmweb --mode upstream:http://127.0.0.1:10809 -s .\mitmproxy_handler.py -p 8899 -q -k
